import sys
import os
import inspect


class CLIError(Exception):
    def __init__(self, fmt, *args):
        super(CLIError, self).__init__(fmt % args if args else str(fmt))

class SwitchNameError(CLIError):
    pass
class PositionalArgError(CLIError):
    pass
class InvalidSwitch(CLIError):
    pass
class InvalidArgument(CLIError):
    pass
class MissingArgument(CLIError):
    pass
class RepeatedSwitch(CLIError):
    pass
class MandatorySwitchMissing(CLIError):
    pass


def switch(*names, **params):
    """
    decorated a function as a CLI switch
    
    == positional arguments ==
    * names - any number of "names" to associate with the switch function. there
      must be at least one name. if a name has already been defined by another 
      switch, if will be an exception, unless the other switch is marked
      as overridable. if the name is one character long, it may be activated
      as a short switch (i.e., "-x"); if the name is two or more characters,
      it can be invoked only as a long switch (i.e., "--xy") 
    
    == keyword arguments ==
    * arg - the "type" of the argument. if not given, no argument is accepted
      for this switch. `type` is any callable object that will convert the 
      string argument into a python object, or raise an exception if the 
      conversion is not possible
    * mandatory - whether or not this switch is mandatory (must be given in 
      the command line). default is False.
    * overridable - whether or not this switch's names are 'weak', i.e., they
      may be overridden by future switches with the same name. if the switch
      is not weak, and another switch attempts to reuse its name, an exception
      will be raised.
    * multiple - whether or not this switch can be repeated in the cli, i.e,
      "-x foo -x bar -x spam"
    * requires - a list of switch names that this switch requires, meaning, in
      order for this switch to be supplied, all the 'required' switches must 
      also have been supplied (i.e., "-x" must be combined with "-y")
    * contradicts - a list of switch names that this switch contradicts, 
      meaning, if this switch has been supplied, all the 'contradicted' 
      switches must not be supplied (i.e., "-x" cannot be combined with "-y")
    * help - the docstring to display in the help context. if not given, will
      default to the docstring of the function. if the function itself has no 
      docstring, the function's name will be used instead.
    """
    if not names:
        raise SwitchNameError("at least one name must be given")
    argtype = params.pop("arg", None)
    mandatory = params.pop("mandatory", False)
    overridable = params.pop("overridable", False)
    multiple = params.pop("multiple", False)
    requires = params.pop("requires", ())
    contradicts = params.pop("contradicts", ())
    help = params.pop("help", None)
    if params:
        raise TypeError("invalid keyword arguments: %r" % (params.keys(),))
    def deco(func):
        def switch_function(self, name, argvalue = None):
            if not multiple and switch_function.cli_info["invoked"]:
                raise RepeatedSwitch("switch %r has already been given", name)
            switch_function.cli_info["invoked"] = True
            if argvalue is None:
                if argtype:
                    raise MissingArgument("switch %r is missing an argument", 
                        name)
                ret = func(self)
            else:
                try:
                    argvalue = argtype(argvalue)
                except Exception, ex:
                    raise InvalidArgument("argument to switch %r: %r", 
                        name, ex)
                ret = func(self, argvalue)
            if ret is not None:
                raise ValueError("switch %r cannot return a value" % (func,))
        
        if argtype:
            named, _, _, defaults = inspect.getargspec(func)
            j = 1 - (len(named) - (len(defaults) if defaults else 0))
            if j >= 0:
                argdefault = defaults[0]
            else:
                argdefault = NotImplemented
            argname = named[1]
        else:
            argname = None
            argdefault = NotImplemented
        
        switch_function.__name__ = func.__name__
        switch_function.__doc__ = func.__doc__
        doc = help
        if doc is None:
            doc = func.__doc__
        if doc is None:
            doc = func.__name__
        switch_function.cli_info = dict(
            names = names,
            doc = " ".join(l.strip() for l in doc.splitlines()),
            argname = argname,
            argtype = argtype,
            argdefault = argdefault,
            invoked = False,
            mandatory = mandatory,
            requires = requires,
            contradicts = contradicts,
            multiple = multiple,
            overridable = overridable,
        )
        return switch_function
    return deco

def mandatory_switch(*names, **params):
    params["mandatory"] = True
    return switch(*names, **params)

def overridable_switch(*names, **params):
    """may be overridden by subclasses with the same switch name"""
    params["overridable"] = True
    return switch(*names, **params)


#===============================================================================
# common CLI types
#===============================================================================
class OneOf(object):
    def __init__(self, *values):
        self.values = dict((v, v) for v in values)
    def __call__(self, value):
        try:
            return self.values[value]
        except KeyError:
            raise TypeError("invalid value %r" % (value,))
    def __str__(self):
        return "one of %r" % (tuple(self.values.keys()),)

class OneOfStr(OneOf):
    def __init__(self, *values):
        self.values = dict((v.lower(), v) for v in values)
    def __call__(self, value):
        return OneOf.__call__(self, value.lower())

class Range(object):
    def __init__(self, type, min, max):
        self.type = type
        self.min = min
        self.max = max
    def __call__(self, value):
        value = self.type(value)
        if value < self.min:
            raise TypeError("value must be >= %r" % (self.min,))
        if value > self.max:
            raise TypeError("value must be <= %r" % (self.max,))
        return value
    def __str__(self):
        return "%s in range [%r..%r]" % (self.type, self.min, self.max)

class InputFile(object):
    def __call__(self, value):
        if not os.path.exists(value):
            raise ValueError("file %r does not exist" % (value,))
        return value
    def __str__(self):
        return "input file"

class OutputFile(object):
    def __init__(self, overwrite = True):
        self.overwrite = overwrite
    def __call__(self, value):
        if not self.overwrite:
            if os.path.exists(value):
                raise ValueError("file %r already exists" % (value,))
        return value
    def __str__(self):
        return "output file"

#===============================================================================
# Main
#===============================================================================
class Main(object):
    """
    This is the baseclass for CLI-based programs. Your program should derive
    from this class and define the `main()` method, which is called as the
    entry point of your program, and the switch methods, which are called to
    process CLI switches.
    
    After processing the command line, `main` will be invoked with the 
    positional arguments passed to this program. If the number of arguments
    passed does not match the number of arguments of `main`, an exception will
    be raised. You can supply default values to the arguments of `main`, in
    which case, if they are not passed, the defaults will take place. Your
    `main` can also accept `*args`, which will allow it to accept any number
    of positional arguments.
    
    Example:
        def main(self) 
            -- no positional arguments may be passed to this program 
        def main(self, first, second):
            -- two positional arguments must be passed to this program
        def main(self, first, second = 8):
            -- one positional argument is required, the second is optional
        def main(self, first, second, *rest):
            -- two positional arguments are required; if more are passed, they
            will be passed as `rest` 
    
    All your executable code should be placed in `main`. You are not intended
    to override __init__ or any other method.
    
    This class defines two overridable built-in switches: verion (-v, --version)
    and help (-h, --help).
    
    Attributes (set these to suit your program):
    * PROGRAM - the program's name, defaults to the name of the class
    * FILENAME - the executable filename, defaults to sys.argv[0]
    * VERSION - the program's version string (no default value)
    """
    PROGRAM = None
    FILENAME = None
    VERSION = "<no version info>"
    
    SHORT_PREFIX = "-"
    LONG_PREFIX = "--"
    END_OF_SWITCHES = "--"
    
    def __init__(self):
        self._cli_switches = {}
        self._invalid_swicthes = []
        for cls in reversed(self.__class__.mro()):
            for name, obj in cls.__dict__.iteritems():
                if hasattr(obj, "cli_info"):
                    self._add_switch(obj)
    
    def _get_switch_func(self, name):
        try:
            return self._cli_switches[name]
        except KeyError:
            return self._handle_invalid_switch(name)
    
    def _handle_invalid_switch(self, name):
        #raise InvalidSwitch(name)
        @switch(name)
        def ignored(self):
            self._invalid_swicthes.append(name)
        return ignored
    
    def _get_arg_value(self, func, argv):
        if func.cli_info["argtype"]:
            try:
                return argv.next()
            except StopIteration:
                return None
        else:
            return None
    
    def _validate(self):
        if self._invalid_swicthes:
            raise InvalidSwitch("unknown switches were passed: %s", 
                ", ".join(repr(n) for n in self._invalid_swicthes))
        for name, func in self._cli_switches.iteritems():
            if func.cli_info["mandatory"] and not func.cli_info["invoked"]:
                raise MandatorySwitchMissing("switch %r is mandatory", name)
            for name2 in func.cli_info["requires"]:
                if not self._cli_switches[name2].cli_info["invoked"]:
                    raise InvalidSwitchCombination("switch %r requires also %r", 
                        name, name2)
            for name2 in func.cli_info["contradicts"]:
                if self._cli_switches[name2].cli_info["invoked"]:
                    raise InvalidSwitchCombination("switch %r contradicts %r", 
                        name, name2)
    
    def _parse(self, argv):
        positional = []
        argv = iter(argv)
        for arg in argv:
            updates = None
            if arg == self.END_OF_SWITCHES:
                positional.extend(argv)
            elif arg.startswith(self.LONG_PREFIX):
                arg = arg[len(self.LONG_PREFIX):]
                if "=" in arg:
                    arg, value = arg.split("=")
                    func = self._get_switch_func(arg)
                else:
                    func = self._get_switch_func(arg)
                    value = self._get_arg_value(func, argv)
                func(self, arg, value)
            elif arg.startswith(self.SHORT_PREFIX):
                arg = arg[len(self.SHORT_PREFIX):]
                for char in arg[:-1]:
                    func = self._get_switch_func(char)
                    func(self, char, None)
                arg = arg[-1]
                func = self._get_switch_func(arg)
                value = self._get_arg_value(func, argv)
                func(self, arg, value)
            else:
                positional.append(arg)
        return positional
    
    def _get_main_argspec(self):
        namedargs, varargs, _, defaults = inspect.getargspec(self.main)
        namedargs = namedargs[1:] # minus 'self'
        defaults = () if defaults is None else defaults
        return namedargs, defaults, varargs
    
    def _preprocess(self, argv):
        positional = self._parse(argv)
        self._validate()
        namedargs, defaults, varargs = self._get_main_argspec()
        required = len(namedargs) - len(defaults)
        if len(positional) < required:
            if defaults:
                raise PositionalArgError("at least %d positional arg(s) "
                    "required, %d given", required, len(positional))
            else:
                raise PositionalArgError("%d positional arg(s) required, "
                    "%d given", required, len(positional))
        if len(positional) > len(namedargs) and not varargs:
            raise PositionalArgError("takes upto %d positional args")
        return positional
    
    def _handle_cli_error(self, ex, argv):
        print "Invocation error! argv = %r" % (argv,)
        print "%s: %s" % (ex.__class__.__name__, ex.message)
        print "Use --help for usage instructions"
    
    def _run(self, argv = None):
        if argv is None:
            argv = sys.argv[1:]
        try:
            positional = self._preprocess(argv)
        except CLIError, ex:
            self._handle_cli_error(ex, argv)
            sys.exit(1)
        retval = self.main(*positional)
        sys.exit(retval)
    
    @classmethod
    def run(cls, argv = None, **kwargs):
        inst = cls(**kwargs)
        return inst._run(argv)
    
    def main(self):
        raise NotImplementedError()
    
    def _add_switch(self, switch_func):
        for name in switch_func.cli_info["names"]:
            if name in self._cli_switches:
                if not self._cli_switches[name].cli_info["overridable"]:
                    raise SwitchNameError("%r is already used by %r", 
                        name, self._cli_switches[name])
            self._cli_switches[name] = switch_func
    
    def switch(self, *names, **params):
        deco = switch(*names, **params)
        def deco2(func):
            func2 = deco(func)
            self._add_switch(func2)
            return func2
        return deco2
    
    #===========================================================================
    # builtin switches (version and help)
    #===========================================================================
    def _print_version_info(self):
        program = self.PROGRAM
        if not program:
            program = self.__class__.__name__
        print "%s version %s" % (program, self.VERSION)
    
    @overridable_switch("h", "help")
    def show_help(self):
        """show this help message and quit"""
        self._print_version_info()
        doc = inspect.getdoc(self)
        if doc:
            print doc.rstrip()
        print
        filename = self.FILENAME
        if not filename:
            filename = os.path.basename(sys.argv[0])
        namedargs, defaults, varargs = self._get_main_argspec()
        arginfo = []
        d = len(namedargs) - len(defaults)
        for i, name in enumerate(namedargs):
            j = i - d
            if j >= 0:
                arginfo.append("%s[= %r]" % (name, defaults[j]))
            else:
                arginfo.append(name)
        if varargs:
            arginfo.append("[%s...]" % (varargs,))
        
        print "Usage: %s [options] %s" % (filename, " ".join(arginfo),)
        print "Options:"
        funcs = {}
        for name, f in self._cli_switches.iteritems():
            if f not in funcs:
                funcs[f] = []
            funcs[f].append(name)
        options = []
        for f, names in funcs.iteritems():
            names.sort()
            opts = []
            for n in names:
                if len(n) == 1:
                    opts.append("%s%s" % (self.SHORT_PREFIX, n))
                else:
                    opts.append("%s%s" % (self.LONG_PREFIX, n))
            opts = ", ".join(opts)
            doc = f.cli_info["doc"].strip()
            if f.cli_info["mandatory"]:
                doc += " (mandatory)"
            if f.cli_info["argtype"]:
                argtype = f.cli_info["argtype"]
                argname = f.cli_info["argname"]
                argdef = f.cli_info["argdefault"]
                opts += " %s" % (argname.upper(),)
                doc += "; %s is %s" % (argname.upper(), argtype)
                if argdef is not NotImplemented:
                    doc += " (defaults to %r)" % (argdef,)
            options.append("  %s %s" % (opts.ljust(30), doc))
        options.sort()
        print "\n".join(options)
        sys.exit()
    
    @overridable_switch("v", "version")
    def show_version(self):
        """show the version info of this program and quit"""
        self._print_version_info()
        sys.exit()


if __name__ == "__main__":
    class MyApp(Main):
        VERSION = "2.2.5"
        
        # defaults
        logfile = None
        quiet = False
        port = 18812
        
        def main(self, a, b, c = 7):
            print "a =", a
            print "b =", b
            print "c =", c
            print "quiet =", self.quiet
            print "logfile =", self.logfile
            print "mode =", self.mode
            print "port =", self.port
        
        @switch("q", "quiet")
        def set_quiet(self):
            self.quiet = True
        
        @switch("logfile", arg=OutputFile())
        def set_logfile(self, filename = logfile):
            """specify a custom logfile"""
            self.logfile = filename
        
        @mandatory_switch("m", "mode", arg=OneOfStr("tcp", "udp"))
        def set_mode(self, mode):
            """set the mode"""
            self.mode = mode
        
        @switch("p", "port", arg=Range(int, 0, 65535))
        def set_port(self, port = port):
            """set the TCP/UDP port"""
            self.port = port    



